{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Importing Brevitas networks into FINN\n",
    "\n",
    "In this notebook we'll go through an example of how to import a Brevitas-trained QNN into FINN. The steps will be as follows:\n",
    "\n",
    "1. Load up the trained PyTorch model\n",
    "2. Call Brevitas FINN-ONNX export and visualize with Netron\n",
    "3. Import into FINN and call cleanup transformations\n",
    "\n",
    "We'll use the following utility functions to print the source code for function calls (`showSrc()`) and to visualize a network using netron (`showInNetron()`) in the Jupyter notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "from finn.util.visualization import showSrc, showInNetron"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load up the trained PyTorch model\n",
    "\n",
    "The FINN Docker image comes with several [example Brevitas networks](https://github.com/maltanar/brevitas_cnv_lfc), and we'll use the LFC-w1a1 model as the example network here. This is a binarized fully connected network trained on the MNIST dataset. Let's start by looking at what the PyTorch network definition looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "class LFC(Module):\n",
      "\n",
      "    def __init__(self, num_classes=10, weight_bit_width=None, act_bit_width=None,\n",
      "                 in_bit_width=None, in_ch=1, in_features=(28, 28), device=\"cpu\"):\n",
      "        super(LFC, self).__init__()\n",
      "        self.device = device\n",
      "\n",
      "        weight_quant_type = get_quant_type(weight_bit_width)\n",
      "        act_quant_type = get_quant_type(act_bit_width)\n",
      "        in_quant_type = get_quant_type(in_bit_width)\n",
      "        stats_op = get_stats_op(weight_quant_type)\n",
      "\n",
      "        self.features = ModuleList()\n",
      "        self.features.append(get_act_quant(in_bit_width, in_quant_type))\n",
      "        self.features.append(Dropout(p=IN_DROPOUT))\n",
      "        in_features = reduce(mul, in_features)\n",
      "        for out_features in FC_OUT_FEATURES:\n",
      "            self.features.append(get_quant_linear(in_features=in_features,\n",
      "                                                  out_features=out_features,\n",
      "                                                  per_out_ch_scaling=INTERMEDIATE_FC_PER_OUT_CH_SCALING,\n",
      "                                                  bit_width=weight_bit_width,\n",
      "                                                  quant_type=weight_quant_type,\n",
      "                                                  stats_op=stats_op))\n",
      "            in_features = out_features\n",
      "            self.features.append(BatchNorm1d(num_features=in_features))\n",
      "            self.features.append(get_act_quant(act_bit_width, act_quant_type))\n",
      "            self.features.append(Dropout(p=HIDDEN_DROPOUT))\n",
      "        self.features.append(get_quant_linear(in_features=in_features,\n",
      "                                   out_features=num_classes,\n",
      "                                   per_out_ch_scaling=LAST_FC_PER_OUT_CH_SCALING,\n",
      "                                   bit_width=weight_bit_width,\n",
      "                                   quant_type=weight_quant_type,\n",
      "                                   stats_op=stats_op))\n",
      "        self.features.append(BatchNorm1d(num_features=num_classes))\n",
      "\n",
      "        for m in self.modules():\n",
      "          if isinstance(m, QuantLinear):\n",
      "            torch.nn.init.uniform_(m.weight.data, -1, 1)\n",
      "\n",
      "    def clip_weights(self, min_val, max_val):\n",
      "        for mod in self.features:\n",
      "            if isinstance(mod, QuantLinear):\n",
      "                mod.weight.data.clamp_(min_val, max_val)\n",
      "    \n",
      "    def forward(self, x):\n",
      "        x = x.view(x.shape[0], -1)\n",
      "        x = 2.0 * x - torch.tensor([1.0]).to(self.device)\n",
      "        for mod in self.features:\n",
      "            x = mod(x)\n",
      "        return x\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from models.LFC import LFC\n",
    "showSrc(LFC)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the network topology is constructed using a few helper functions that generate the quantized linear layers and quantized activations. The bitwidth of the layers is actually parametrized in the constructor, so let's instantiate a 1-bit weights and activations version of this network. We also have pretrained weights for this network, which we will load into the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LFC(\n",
       "  (features): ModuleList(\n",
       "    (0): QuantHardTanh(\n",
       "      (act_quant_proxy): ActivationQuantProxy(\n",
       "        (fused_activation_quant_proxy): FusedActivationQuantProxy(\n",
       "          (activation_impl): Identity()\n",
       "          (tensor_quant): ClampedBinaryQuant(\n",
       "            (scaling_impl): StandaloneScaling(\n",
       "              (restrict_value): RestrictValue(\n",
       "                (forward_impl): Sequential(\n",
       "                  (0): PowerOfTwo()\n",
       "                  (1): ClampMin()\n",
       "                )\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (1): Dropout(p=0.2)\n",
       "    (2): QuantLinear(\n",
       "      in_features=784, out_features=1024, bias=False\n",
       "      (weight_reg): WeightReg()\n",
       "      (weight_quant): WeightQuantProxy(\n",
       "        (tensor_quant): BinaryQuant(\n",
       "          (scaling_impl): StandaloneScaling(\n",
       "            (restrict_value): RestrictValue(\n",
       "              (forward_impl): Sequential(\n",
       "                (0): PowerOfTwo()\n",
       "                (1): Identity()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (bias_quant): BiasQuantProxy()\n",
       "    )\n",
       "    (3): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (4): QuantHardTanh(\n",
       "      (act_quant_proxy): ActivationQuantProxy(\n",
       "        (fused_activation_quant_proxy): FusedActivationQuantProxy(\n",
       "          (activation_impl): Identity()\n",
       "          (tensor_quant): ClampedBinaryQuant(\n",
       "            (scaling_impl): StandaloneScaling(\n",
       "              (restrict_value): RestrictValue(\n",
       "                (forward_impl): Sequential(\n",
       "                  (0): PowerOfTwo()\n",
       "                  (1): ClampMin()\n",
       "                )\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (5): Dropout(p=0.2)\n",
       "    (6): QuantLinear(\n",
       "      in_features=1024, out_features=1024, bias=False\n",
       "      (weight_reg): WeightReg()\n",
       "      (weight_quant): WeightQuantProxy(\n",
       "        (tensor_quant): BinaryQuant(\n",
       "          (scaling_impl): StandaloneScaling(\n",
       "            (restrict_value): RestrictValue(\n",
       "              (forward_impl): Sequential(\n",
       "                (0): PowerOfTwo()\n",
       "                (1): Identity()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (bias_quant): BiasQuantProxy()\n",
       "    )\n",
       "    (7): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (8): QuantHardTanh(\n",
       "      (act_quant_proxy): ActivationQuantProxy(\n",
       "        (fused_activation_quant_proxy): FusedActivationQuantProxy(\n",
       "          (activation_impl): Identity()\n",
       "          (tensor_quant): ClampedBinaryQuant(\n",
       "            (scaling_impl): StandaloneScaling(\n",
       "              (restrict_value): RestrictValue(\n",
       "                (forward_impl): Sequential(\n",
       "                  (0): PowerOfTwo()\n",
       "                  (1): ClampMin()\n",
       "                )\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (9): Dropout(p=0.2)\n",
       "    (10): QuantLinear(\n",
       "      in_features=1024, out_features=1024, bias=False\n",
       "      (weight_reg): WeightReg()\n",
       "      (weight_quant): WeightQuantProxy(\n",
       "        (tensor_quant): BinaryQuant(\n",
       "          (scaling_impl): StandaloneScaling(\n",
       "            (restrict_value): RestrictValue(\n",
       "              (forward_impl): Sequential(\n",
       "                (0): PowerOfTwo()\n",
       "                (1): Identity()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (bias_quant): BiasQuantProxy()\n",
       "    )\n",
       "    (11): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (12): QuantHardTanh(\n",
       "      (act_quant_proxy): ActivationQuantProxy(\n",
       "        (fused_activation_quant_proxy): FusedActivationQuantProxy(\n",
       "          (activation_impl): Identity()\n",
       "          (tensor_quant): ClampedBinaryQuant(\n",
       "            (scaling_impl): StandaloneScaling(\n",
       "              (restrict_value): RestrictValue(\n",
       "                (forward_impl): Sequential(\n",
       "                  (0): PowerOfTwo()\n",
       "                  (1): ClampMin()\n",
       "                )\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (13): Dropout(p=0.2)\n",
       "    (14): QuantLinear(\n",
       "      in_features=1024, out_features=10, bias=False\n",
       "      (weight_reg): WeightReg()\n",
       "      (weight_quant): WeightQuantProxy(\n",
       "        (tensor_quant): BinaryQuant(\n",
       "          (scaling_impl): StandaloneScaling(\n",
       "            (restrict_value): RestrictValue(\n",
       "              (forward_impl): Sequential(\n",
       "                (0): PowerOfTwo()\n",
       "                (1): Identity()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (bias_quant): BiasQuantProxy()\n",
       "    )\n",
       "    (15): BatchNorm1d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "trained_lfc_w1a1_checkpoint = \"/workspace/brevitas_cnv_lfc/pretrained_models/LFC_1W1A/checkpoints/best.tar\"\n",
    "lfc = LFC(weight_bit_width=1, act_bit_width=1, in_bit_width=1).eval()\n",
    "checkpoint = torch.load(trained_lfc_w1a1_checkpoint, map_location=\"cpu\")\n",
    "lfc.load_state_dict(checkpoint[\"state_dict\"])\n",
    "lfc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have now instantiated our trained PyTorch network. Let's try to run an example MNIST image through the network using PyTorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAARX0lEQVR4nO3dfYyVZXrH8d/FoDAw8iYRCaisG/5QqmUbgk1KyOKmxlUMbKJm/aPauAmarMmqTVqz/UOSaqJVa/pH3YStL9CsmiWoq0a7a82mWo1GNFQQW1CULGR4E5H3t+HqH/NgZ3We6549z3nOc9z7+0kmM3Ouec65OTM/zsv13Pdt7i4Af/xGNT0AAJ1B2IFMEHYgE4QdyARhBzIxupM3Zma89Z+ZUaPKH09OnTpV23VXvf6enp6wPjAw0PJ1183dbbjLK4XdzK6U9M+SeiT9q7vfV+X6cmU27O/mS6k/6ip/eKNHx38CqcCk6r29vaW1Q4cOhcem9PX1hfUDBw6U1lIt50mTJoX1zz77LKx3o5afxptZj6R/kfR9SRdLusHMLm7XwAC0V5XX7PMlfeTuW9z9uKSnJS1pz7AAtFuVsM+Q9Lsh328rLvs9ZrbMzNaa2doKtwWgotrfoHP3FZJWSLxBBzSpyiP7dknnDfl+ZnEZgC5UJezvSJptZt8yszMl/VDS8+0ZFoB2a/lpvLufNLPbJP1ag623x9z9g7aNLCPjx48P6wcPHmz5useMGRPWjx07FtZTbcFx48aF9ai9lmoppqSOj9prqT76vn37WhpTN6v0mt3dX5L0UpvGAqBGnC4LZIKwA5kg7EAmCDuQCcIOZIKwA5mwTq4um+vpsqled6qXffTo0bA+duzYlo9Nia676vWfffbZYb3qNNLofp06dWp47O7du8N6amrwyZMnw3qdyuaz88gOZIKwA5kg7EAmCDuQCcIOZIKwA5mg9fYNkGrNVfkd1nnddUtNDa6yem1q6m5qanCTS03TegMyR9iBTBB2IBOEHcgEYQcyQdiBTBB2IBP02TvgrLPOCuvRbqOSNHHixLB+4sSJ0lpqN9LUFNbPP/88rC9YsCCs33rrraW1VC/6jjvuCOtbt24N601OM20SfXYgc4QdyARhBzJB2IFMEHYgE4QdyARhBzJBn/0b4JFHHgnrUS871Wuuuox1b29vWI+ktk2+5JJLwvqmTZvC+vHjx0trZ5xxRnhsdO6ClP53HzlyJKzXqazPXmnLZjP7VNIBSQOSTrr7vCrXB6A+lcJeWOTue9pwPQBqxGt2IBNVw+6SfmNm75rZsuF+wMyWmdlaM1tb8bYAVFD1afwCd99uZudIesXM/sfdXxv6A+6+QtIKiTfogCZVemR39+3F512SnpU0vx2DAtB+LYfdzMab2Vmnv5Z0haQN7RoYgPaq8jR+mqRniz7taElPuvu/t2VUf2RSWzYvWrQorF922WVhPeqVHzx4MDw21W/u6+sL66nzNKI566m11x999NGWr1uS7rzzztLaW2+9FR5b93bSTWg57O6+RdKftnEsAGpE6w3IBGEHMkHYgUwQdiAThB3IBFNcu0Bqqubs2bPD+v79+0trEyZMCI+NpoFK6SmwVbZ8TrX9UlJLcO/du7e0tnTp0vDYdevWhfVUSzLV8qwTS0kDmSPsQCYIO5AJwg5kgrADmSDsQCYIO5CJdiw42TFRT7fOfnBK6thU/ZZbbgnrq1atCuszZ85s+bZTffZ77rknrK9evTqsn3nmmaW1K664Ijz2wQcfDOuprbCj2168eHF47LZt28L6nj3fvDVWeWQHMkHYgUwQdiAThB3IBGEHMkHYgUwQdiATHZ/Pnup3Rzo51naqOvd54cKFYf2iiy4qrY0bNy48dvTo+FSLNWvWhPUtW7aE9SpSyz3PmTMnrKfu90jq75T57AC6FmEHMkHYgUwQdiAThB3IBGEHMkHYgUx0vM8+alT5/y9V54XXqcpc+lOnTlW67eg+S9VPnjwZHjt+/PiwfujQobCe2o46+p2l5tJfffXVYf3pp58O61X67Kk17VP3a5Na7rOb2WNmtsvMNgy5bIqZvWJmm4vPk9s5WADtN5Kn8U9IuvIrl90l6VV3ny3p1eJ7AF0sGXZ3f03SV/fRWSJpZfH1SknxXjoAGtfqGnTT3L2/+HqHpGllP2hmyyQta/F2ALRJ5QUn3d2jDRvdfYWkFRIbOwJNarX1ttPMpktS8XlX+4YEoA6thv15STcVX98k6VftGQ6AuiT77Gb2lKTvSpoqaaekuyU9J+mXks6XtFXS9e5evhn2/19XbU/jq64bX7UeSfVkU3uoR/uvV9Xb2xvWjxw5EtZT5wBUOcfgwgsvDOsff/xxy9edGldqTfqUw4cPVzq+irI+e/I1u7vfUFL6XqURAegoTpcFMkHYgUwQdiAThB3IBGEHMsGWzYVUC3JgYCCsR3p6esJ61WWHozZRqsWUmsKakrr+aNvkqCZJixYtamlMp0W/0xMnToTHpqa4Vvl7aAqP7EAmCDuQCcIOZIKwA5kg7EAmCDuQCcIOZKKr+ux1budcdTnnKuq+7QMHDpTWUv3iVK87dXyqTx8tF51axvq6664L60ePHg3rY8eOLa2l+uyp31mTWzK3ikd2IBOEHcgEYQcyQdiBTBB2IBOEHcgEYQcy0fE+ezS3u5t75dGSyanllFPq3Fb50ksvDY+dM2dOWE8tJf3cc8+F9UjUB5ekhQsXhvUqW3inlqGOzl2Qqi/B3QQe2YFMEHYgE4QdyARhBzJB2IFMEHYgE4QdyETH++zRnPU6++ipufKped1RT3j06PhuXLp0aVhPHb9kyZKwPmbMmNLa3Llzw2MnTZoU1lO97Ndff73l42fPnh0em1qbPdXrXr9+fWnt8ssvD4+N7lOpO/voKclHdjN7zMx2mdmGIZctN7PtZrau+Liq3mECqGokT+OfkHTlMJc/7O5zi4+X2jssAO2WDLu7vyZpbwfGAqBGVd6gu83M3i+e5k8u+yEzW2Zma81sbYXbAlBRq2H/maRvS5orqV/SQ2U/6O4r3H2eu89r8bYAtEFLYXf3ne4+4O6nJP1c0vz2DgtAu7UUdjObPuTbH0jaUPazALqDpfqoZvaUpO9Kmippp6S7i+/nSnJJn0q6xd37kzdmFt5Yqt+cmvcdmTVrVli/5pprwvrixYtLa6l516l526m509H+61K8hnlfX194bErVed3R7/SLL74Ij504cWJYT9m8eXNpbdWqVeGxDz1U+spUUnf32d192JNKkifVuPsNw1z8aOURAegoTpcFMkHYgUwQdiAThB3IBGEHMpFsvbX1xsw8Wna5zimud999d1hfvnx5WN+zZ09pberUqa0M6UuprYf37o2nJkT1Cy64IDw21RZMbdmccuzYsdJaahpp6u8h1YqNpi2ntlx++eWXw/rNN98c1pvc0rms9cYjO5AJwg5kgrADmSDsQCYIO5AJwg5kgrADmeh4nz2qV9maODXVMtX3rLLt8q5du8L61q1bw/oDDzwQ1levXh3W580rXwTo4YcfDo9Nbdk8eXLpimOSpG3btoX16Hf6xBNPhMd+8sknYf3aa68N69HU46rTa1988cWwnpoyXSf67EDmCDuQCcIOZIKwA5kg7EAmCDuQCcIOZKKjffZRo0Z5ND/6+PHj4fHnnHNOaW337t3hsak+e2rudNQvTm0HvWnTprA+ZcqUsJ5atjha7vn8888Pj03NZ08t771v376wfuONN5bWXnjhhfDYlNQ6AtFy0YsWLQqPTa0xkLpfUst/14k+O5A5wg5kgrADmSDsQCYIO5AJwg5kgrADmeiq+exVpPqeK1euDOvXX399y9d/+PDh8Nhx48aF9dS2yKl5/gMDA6W11Lrvb775Zlh/8sknw/q6devC+htvvFFaS51fkOrhp37n0Xkb8+fPD499++23w/rjjz8e1lPrytep5T67mZ1nZr81s41m9oGZ/aS4fIqZvWJmm4vP8SoHABo1kqfxJyX9jbtfLOnPJf3YzC6WdJekV919tqRXi+8BdKlk2N29393fK74+IOlDSTMkLZF0+rnxSklL6xokgOriFz1fYWazJH1H0tuSprl7f1HaIWlayTHLJC1rfYgA2mHE78abWZ+kNZJud/f9Q2s++C7fsG++ufsKd5/n7uWrIgKo3YjCbmZnaDDov3D3Z4qLd5rZ9KI+XVK8xCqARiVbbzY4f3OlpL3ufvuQyx+Q9Jm732dmd0ma4u5/m7iu8MbOPffccCw7duwI65Fo+15JmjlzZli/9957S2szZswIj01tuZzaujjaLlqS7r///tLaxo0bw2NTU1xT2yKnpKYtR1JtwxMnToT1aOpx6u9+woQJYb3qlOk6lbXeRvKa/S8k/ZWk9WZ2uqn6U0n3Sfqlmf1I0lZJcaMaQKOSYXf3/5JU9l/k99o7HAB14XRZIBOEHcgEYQcyQdiBTBB2IBMdneLa09PjUV83NVU06n3u37+/tCZJfX19YT3VN416vlX6vVK655s6RyDqZad6+MeOHQvrVUW/79Ryzampwam/lyq/s5SqY6sTS0kDmSPsQCYIO5AJwg5kgrADmSDsQCYIO5CJrlpKOjWHOOqlp5YVrjove/r06aW1/v7+0tpI9Pb2hvXUls11XndqGetDhw6F9SpzylNGjYofq6rMKW/6/IQq6LMDmSPsQCYIO5AJwg5kgrADmSDsQCYIO5CJruqzA6iOPjuQOcIOZIKwA5kg7EAmCDuQCcIOZIKwA5lIht3MzjOz35rZRjP7wMx+Uly+3My2m9m64uOq+ocLoFXJk2rMbLqk6e7+npmdJeldSUs1uB/7QXd/cMQ3xkk1QO3KTqoZyf7s/ZL6i68PmNmHkma0d3gA6vYHvWY3s1mSviPp7eKi28zsfTN7zMwmlxyzzMzWmtnaSiMFUMmIz403sz5J/ynpXnd/xsymSdojySX9gwaf6t+cuA6exgM1K3saP6Kwm9kZkl6U9Gt3/6dh6rMkvejuf5K4HsIO1KzliTA2uDzoo5I+HBr04o27034gaUPVQQKoz0jejV8g6XVJ6yWdXpv3p5JukDRXg0/jP5V0S/FmXnRdPLIDNav0NL5dCDtQP+azA5kj7EAmCDuQCcIOZIKwA5kg7EAmCDuQCcIOZIKwA5kg7EAmCDuQCcIOZIKwA5kg7EAmkgtOttkeSVuHfD+1uKwbdevYunVcEmNrVTvHdkFZoaPz2b9242Zr3X1eYwMIdOvYunVcEmNrVafGxtN4IBOEHchE02Ff0fDtR7p1bN06LomxtaojY2v0NTuAzmn6kR1AhxB2IBONhN3MrjSz/zWzj8zsribGUMbMPjWz9cU21I3uT1fsobfLzDYMuWyKmb1iZpuLz8PusdfQ2LpiG+9gm/FG77umtz/v+Gt2M+uRtEnSX0raJukdSTe4+8aODqSEmX0qaZ67N34ChpktlHRQ0qrTW2uZ2T9K2uvu9xX/UU5297/rkrEt1x+4jXdNYyvbZvyv1eB9187tz1vRxCP7fEkfufsWdz8u6WlJSxoYR9dz99ck7f3KxUskrSy+XqnBP5aOKxlbV3D3fnd/r/j6gKTT24w3et8F4+qIJsI+Q9Lvhny/Td2137tL+o2ZvWtmy5oezDCmDdlma4ekaU0OZhjJbbw76SvbjHfNfdfK9udV8Qbd1y1w9z+T9H1JPy6ernYlH3wN1k29059J+rYG9wDsl/RQk4MpthlfI+l2d98/tNbkfTfMuDpyvzUR9u2Szhvy/czisq7g7tuLz7skPavBlx3dZOfpHXSLz7saHs+X3H2nuw+4+ylJP1eD912xzfgaSb9w92eKixu/74YbV6futybC/o6k2Wb2LTM7U9IPJT3fwDi+xszGF2+cyMzGS7pC3bcV9fOSbiq+vknSrxocy+/plm28y7YZV8P3XePbn7t7xz8kXaXBd+Q/lvT3TYyhZFwXSvrv4uODpscm6SkNPq07ocH3Nn4k6WxJr0raLOk/JE3porH9mwa39n5fg8Ga3tDYFmjwKfr7ktYVH1c1fd8F4+rI/cbpskAmeIMOyARhBzJB2IFMEHYgE4QdyARhBzJB2IFM/B+tIjCppYWKvAAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from pkgutil import get_data\n",
    "import onnx\n",
    "import onnx.numpy_helper as nph\n",
    "raw_i = get_data(\"finn\", \"data/onnx/mnist-conv/test_data_set_0/input_0.pb\")\n",
    "input_tensor = onnx.load_tensor_from_string(raw_i)\n",
    "input_tensor_npy = nph.to_array(input_tensor)\n",
    "input_tensor_pyt = torch.from_numpy(input_tensor_npy).float()\n",
    "imgplot = plt.imshow(input_tensor_npy.reshape(28,28), cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0602, 0.0147, 0.5844, 0.0445, 0.0270, 0.0185, 0.0595, 0.0082, 0.1689,\n",
       "        0.0141])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch.nn.functional import softmax\n",
    "# do forward pass in PyTorch/Brevitas\n",
    "produced = lfc.forward(input_tensor_pyt).detach()\n",
    "probabilities = softmax(produced, dim=-1).flatten()\n",
    "probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYIAAAEICAYAAABS0fM3AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAclUlEQVR4nO3debgdVZ3u8e9LIDIFgiQqJEDCaMcJ8RgQFZGhO3QreBERrmjHK9NtwkXx2oIitrRTO+GErUAQlEZuiGhHjQRolBYVSEAEkoCEEEgYD5Mg+gCB9/5RdXBz2GefnaHq5Jx6P8+zn1PDqlq/vXeyf7VWVa2SbSIiornWG+oAIiJiaCURREQ0XBJBRETDJRFERDRcEkFERMMlEURENFwSQUSXJFnSjuX0tyV9YjX38ydJ26/d6DrWJ0nflfSIpGvrqjeGjySCEUjSMkn7tVm+t6Rnyx+ivtdPWtbvLOkiSQ9K+qOkGyWdKGnUGsYzQ9ICSU9KOncVtz1U0m8k/VnSLwcp2/r+Hpd0q6T3r0nsA7F9rO1/HaycpF9KOrLftpvaXlpFXAN4E7A/MNH21DXdmaRJZVJcf81Di3VBEkHz3FP+EPW93g4gaQfgGmA58CrbmwPvAnqAMWtaJ/Bp4JzV2PZh4KvA57uty/amwGbAR4GzJE3pX6hhP2LbActsP7GqGzbsc2qsJILo8yngN7ZPtH0vgO1bbf9P24/2LyzprZJuapm/TNL8lvlfSXpHuZ+Lbf8YeKjNfraQ9FNJvWXXxU8lTexbb/ty27MokknXXPgx8AgwRdJ0Sb+WdLqkh4B/kfQiSV+SdJek+8vuno1aYvuIpHsl3SPpf/WL+1xJn26ZP0jSDZIek3S7pGmSPgO8Gfhm2Ur5Zlm2tYtpc0nfK9//nZJOkbReuW66pKvKGB+RdIekA1rqnC5padn6uUPSe9p8vh8AzgbeUMbwqXL5UZKWSHpY0hxJW7dsY0nHSboNuG2wz7r8LL4l6edlHb+W9DJJXy3jvkXSa1vKn1R+Ro9LWiTpf7SsGyXpy2Wr9I6yNflc66P8vGaW38vdkj69pi3WSCKIv9oPmL0K5a8GdpI0TtIGwKuBrSWNKX9Me4BfdbGf9YDvUhy1bgv8BfjmKkXehqT1yh+YsUBfwtodWAq8FPgMRStjZ2BXYEdgAnBquf004P9SdKnsRPH5DFTXVOB7wEfK+vaiOAL/OMVnMKNsfc1os/k3gM2B7YG3AO8DWruzdgduBcYBXwBmqrAJ8HXgANtjgD2BG/rv3PZM4Fjgt2UMn5S0D/A54FBgK+BO4MJ+m76jrPsFrakBHAqcUsb5JPBb4PpyfjbwlZayt1MkyM0pDkDOl7RVue4o4ACK72S3Mo5W5wIrKb6v1wJ/CxxJrBnbeY2wF7AM2K/N8r2BZ4FHW16HluueBqatYj2/Ag4G9gAuBWYB04C3Aje2Kf9p4NxB9rkr8Eib5UcCvxxk29b39zDFD+Nh5brpwF0tZQU8AezQsuwNwB3l9DnA51vW7QwY2LGcPxf4dDn9HeD0AWL6JXBkv2Wm+CEbBTwFTGlZd0zf+yxjXtKybuNy25cBm5Tv853ARoN8LtOBq1rmZwJfaJnftPz+J7XEt0+H/U0qy6zf8lmc1bL+eGBxy/yrgEc77O8G4KBy+grgmJZ1+/XVRZHAn2x9v8DhwC+G8v/bSHil/6957rE9sc3yhyiODtuS9G3giHL2s7Y/C1xJ8eO7opx+hOKo9slyflCSNgZOp0ggW5SLx0gaZfuZbvbRz0DvD4rzH33GU/ywXifpuXAofpwBtgauayl/Z4c6twHmrnqojAM26LfvOylaJn3u65uw/ecy1k1t3yfp3RStlpmSfg182PYtXdS7NcXRet9+/1R2l02gOIiA539W3bi/ZfovbeY37ZuR9D7gRIqEQrluXEtsrXW3Tm9H8Xnd2/KdrbcasUY/6RqKPpdTHF225eIqmb4TzJ8tF/clgr3K6SspEsFb6DIRAB8GdgF2t71ZuS8ofpTXttahdh+k+IF6he2x5WtzFyeaAe6l+IHvs22H/S4Hduiizv4epDgS365fPXd32OavO7bn2d6fIoHfApzVzXYU51ueq7PsZtqyX72VDEssaTuKOGcAW9oeC9zMX7/ve4HWRN76HSynOMgY1/KdbWb7FVXE2iRJBCPXBpI2bHkN1vr7JLCnpC9KehmApB0lnS9p7ADb/IbiR3wqcK3thRQ/MLsD/91XSNL6kjakONoe1S+eMRQ/yI9KenEZBy3bjiq3XR9Yr9x2g1X4HNqy/SzFD9Lpkl5S1jVB0t+VRWYB0yVNKVstnxxgV1B0tbxf0r7luYkJkl5errufov+/XQzPlPV8pjy3sh3FkfL5g8Uv6aXlCepNKH4c/0TRLdaNH5Tx7irpRcBngWtsL+ty+zWxCUWS6QVQcXnvK1vWzwJOKD/DsRRXfgHg4iKGS4EvS9qs/Kx3kPSWGuIe0ZIIRq65FD+wfa9/6VTY9u0UfeSTgIWS/gj8EFgAPD7ANk9QdDEstP1Uufi3wJ22H2gpekoZw0kU3Ut/KZdBcWnoRhRHx1cDl/Sr5r1l+X+nOMH4F7o/8h3MR4ElwNWSHqNoFe1Svrefl7FdUZa5YqCd2L6W4gTv6cAfKVpDfUfcXwMOKa+e+XqbzY+nOFexFLgKuIDuLrNdjyJp3ENxPuQtwP/uYjtsXw58guL7vZeiNXNYN9uuKduLgC9T/Du5n+L8wa9bipxF8WN/I/A7in/HK4G+bsL3AaOBRRRdkbPp0KUZ3VF5wiUiYp1TXi77bdvbDVo4VltaBBGxzpC0kaS/L7sTJ1B0yf1oqOMa6dIiiIh1Rnk+5krg5RTdgD8DTrD92JAGNsIlEURENFy6hiIiGm7Y3VA2btw4T5o0aajDiIgYVq677roHbY9vt27YJYJJkyaxYMGCoQ4jImJYkTTg3fGVdg2pGIHx1nKUw5MGKHNoOQLhQkkXVBlPRES8UGUtgnJo2DMoRm9cAcyXNKe8oaSvzE7AycAbbT/Sd4dnRETUp8oWwVSKkROXlnedXggc1K/MUcAZth8B6Hc3akRE1KDKRDCB548KuILnj6oIxdC+O5cPsri6HAP+BSQdreJRhwt6e3srCjciopmG+vLR9Ske+rE3xbjiZ7Ub4Mz2mbZ7bPeMH9/2pHdERKymKhPB3Tx/CNmJvHB43RXAHNtP274D+ANFYoiIiJpUmQjmUzzKcLKk0RSjG87pV+bHFK0BJI2j6CpaWmFMERHRT2WJwPZKiodPzAMWA7NsL5R0mqQDy2LzgIckLQJ+AXzE9gsecB4REdUZdmMN9fT0ODeURUSsGknX2e5pt27Y3Vkcq+70y/5QeR0f2n/nyuuIiGoM9VVDERExxJIIIiIaLokgIqLhkggiIhouiSAiouGSCCIiGi6JICKi4ZIIIiIaLokgIqLhkggiIhouiSAiouGSCCIiGi6JICKi4ZIIIiIaLokgIqLhkggiIhouiSAiouGSCCIiGi6JICKi4ZIIIiIaLokgIqLhkggiIhouiSAiouGSCCIiGi6JICKi4SpNBJKmSbpV0hJJJ7VZP11Sr6QbyteRVcYTEREvtH5VO5Y0CjgD2B9YAcyXNMf2on5F/5/tGVXFERERnVXZIpgKLLG91PZTwIXAQRXWFxERq6HKRDABWN4yv6Jc1t87Jd0oabakbSqMJyIi2hjqk8U/ASbZfjVwGXBeu0KSjpa0QNKC3t7eWgOMiBjpqkwEdwOtR/gTy2XPsf2Q7SfL2bOB17Xbke0zbffY7hk/fnwlwUZENFWViWA+sJOkyZJGA4cBc1oLSNqqZfZAYHGF8URERBuVXTVke6WkGcA8YBRwju2Fkk4DFtieA/wfSQcCK4GHgelVxRMREe1VlggAbM8F5vZbdmrL9MnAyVXGEBERnQ31yeKIiBhiSQQREQ2XRBAR0XBJBBERDZdEEBHRcEkEERENl0QQEdFwSQQREQ2XRBAR0XBJBBERDZdEEBHRcEkEERENl0QQEdFwgyYCScdL2qKOYCIion7dtAheCsyXNEvSNEmqOqiIiKjPoInA9inATsBMigfH3Cbps5J2qDi2iIioQVfnCGwbuK98rQS2AGZL+kKFsUVERA0GfUKZpBOA9wEPUjxg/iO2n5a0HnAb8M/VhhgREVXq5lGVLwYOtn1n60Lbz0p6WzVhRUREXbrpGtq+fxKQ9H0A24sriSoiImrTTSJ4ReuMpFHA66oJJyIi6jZgIpB0sqTHgVdLeqx8PQ48APxnbRFGRESlBkwEtj9newzwRdubla8xtre0fXKNMUZERIUGPFks6eW2bwEukrRb//W2r680soiIqEWnq4Y+DBwFfLnNOgP7VBJRRETUasBEYPuo8u9b6wsnIiLq1qlr6OBOG9q+eO2HExERdevUNfT2DusMDJoIJE0DvgaMAs62/fkByr0TmA283vaCwfYbERFrT6euofevyY7L+w3OAPYHVlCMYDrH9qJ+5cYAJwDXrEl9ERGxejp1DR1h+3xJJ7Zbb/srg+x7KrDE9tJyfxcCBwGL+pX7V+DfgI90HXVERKw1ne4s3qT8O2aA12AmAMtb5leUy55TXpa6je2fddqRpKMlLZC0oLe3t4uqIyKiW526hr5T/v1UFRWXo5d+heIZBx3ZPhM4E6Cnp8dVxBMR0VTdPKpye0k/kdQr6QFJ/ylp+y72fTewTcv8xHJZnzHAK4FfSloG7AHMkdTTffgREbGmuhl07gJgFrAVsDVwEfCDLrabD+wkabKk0cBhwJy+lbb/aHuc7Um2JwFXAwfmqqGIiHp1kwg2tv192yvL1/nAhoNtZHslMAOYBywGZtleKOk0SQeuWdgREbG2dLpq6MXl5M8lnQRcSHH/wLuBud3s3Pbc/mVtnzpA2b272WdERKxdnW4ou47ih1/l/DEt6wxkBNKIiBGg01VDk+sMJCIihkY3zyxG0iuBKbScG7D9vaqCioiI+gyaCCR9EtibIhHMBQ4ArgKSCCIiRoBurho6BNgXuK8cf+g1wOaVRhUREbXpJhH8xfazwEpJm1E8s3ibQbaJiIhhoptzBAskjQXOoriS6E/AbyuNKiIiajNoIrD9T+XktyVdAmxm+8Zqw4qIiLp0e9XQwcCbKO4fuApIIoiIGCG6GXTuW8CxwE3AzcAxks6oOrCIiKhHNy2CfYC/sW0ASecBCyuNKiIiatPNVUNLgG1b5rcpl0VExAjQadC5n1CcExgDLJZ0bblqKnDtQNtFRMTw0qlr6Eu1RREREUOm06BzV/ZNS3op8Ppy9lrbD1QdWERE1KObq4YOpegKehdwKHCNpEOqDiwiIurRzVVDHwde39cKkDQeuByYXWVgERFRj26uGlqvX1fQQ11uFxERw0A3LYJLJM3jrw+s7/pRlRERse7rmAgkCfg6xYniN5WLz7T9o6oDi4iIenRMBLYtaa7tVwEX1xRTRETUqJu+/uslvX7wYhERMRx1c45gd+AIScuAJwBRNBZeXWVgERFRj24Swd9VHkVERAyZTmMNvQT4GLAjxRDUn7P9WF2BRUREPTqdI/geRVfQN4BNKa4eioiIEaZTItjK9sdtz7N9PLDK5wQkTZN0q6Qlkk5qs/5YSTdJukHSVZKmrGodERGxZjpeNSRpC0kvlvRiYFS/+Y4kjQLOAA4ApgCHt/mhv8D2q2zvCnwB+MrqvY2IiFhdnU4Wbw5cR3GVUJ/ry78Gth9k31OBJbaXAki6EDgIWNRXoN85h03K/UZERI06DUM9aQ33PQFY3jK/guJS1OeRdBxwIjCa4rGYLyDpaOBogG233bZdkYiIWE1DPnic7TNs7wB8FDhlgDJn2u6x3TN+/Ph6A4yIGOGqTAR3UzzfuM/EctlALgTeUWE8ERHRRpWJYD6wk6TJkkYDhwFzWgtI2qll9h+A2yqMJyIi2uh0Q1nHK4NsPzzI+pWSZgDzgFHAObYXSjoNWGB7DjBD0n7A08AjwD+u6huIiIg10+mqoesoruIRsC3FD7WAscBdwOTBdm57Lv2eXWD71JbpE1Y95IiIWJsG7BqyPdn29hSPpXy77XG2twTeBlxaV4AREVGtbs4R7FEe2QNg++fAntWFFBERdepm9NF7JJ0CnF/Ovwe4p7qQIiKiTt20CA4HxgM/onhK2fhyWUREjACDtgjKq4NOkLSJ7SdqiCkiImo0aItA0p6SFgGLy/nXSPpW5ZFFREQtuukaOp3iKWUPAdj+PbBXlUFFRER9urqz2PbyfoueqSCWiIgYAt1cNbRc0p6AJW0AnEDZTRQREcNfNy2CY4HjKIaVvhvYFfinKoOKiIj6dNMi2MX2e1oXSHoj8OtqQoqIiDp10yL4RpfLIiJiGOo0+ugbKIaSGC/pxJZVm1GMJhoRESNAp66h0cCmZZkxLcsfAw6pMqiIiKhPp2cWXwlcKelc23fWGFNERNSom3MEZ0sa2zcjaQtJ8yqMKSIiatRNIhhn+9G+GduPAC+pLqSIiKhTN4ngWUnb9s1I2o7iyWURETECdHMfwceBqyRdSfGoyjcDR1caVURE1KabYagvkbQbsEe56IO2H6w2rIiIqMuAXUOSXl7+3Y3i4fX3lK9ty2URETECdGoRfBg4Cvhym3UG9qkkooiIqFWn+wiOKv++tb5wIiKibp2GmDi404a2L1774URERN06dQ29vfz7Eooxh64o598K/IbiQfYRETHMdeoaej+ApEuBKbbvLee3As6tJbqIiKhcNzeUbdOXBEr3U1xFFBERI0A3ieC/JM2TNF3SdOBnwOXd7FzSNEm3Sloi6aQ260+UtEjSjZL+q7xrOSIiajRoIrA9A/g28Jrydabt4wfbTtIo4AzgAGAKcLikKf2K/Q7osf1qYDbwhVULPyIi1lQ3Q0wAXA88bvtySRtLGmP78UG2mQossb0UQNKFwEHAor4Ctn/RUv5q4IjuQ4+IiLVh0BaBpKMojta/Uy6aAPy4i31PAJa3zK8olw3kA8DPB4jhaEkLJC3o7e3touqIiOhWN+cIjgPeSPFkMmzfxloehlrSEUAP8MV2622fabvHds/48ePXZtUREY3XTdfQk7afkgSApPXpbhjqu4FtWuYnlsueR9J+FCOcvsX2k13sNyIi1qJuWgRXSvoYsJGk/YGLgJ90sd18YCdJkyWNBg4D5rQWkPRaii6nA20/sGqhR0TE2tBNIvgo0AvcBBwDzAVOGWwj2yuBGcA8YDEwy/ZCSadJOrAs9kVgU+AiSTdImjPA7iIioiIdu4bKS0AX2n45cNaq7tz2XIrE0brs1Jbp/VZ1nxERsXZ1bBHYfga4tfVRlRERMbJ0c7J4C2ChpGuBJ/oW2j5w4E0iImK46CYRfKLyKCIiYsh0eh7BhsCxwI4UJ4pnlieAIyJiBOl0juA8ipu8bqIYL6jdIysjImKY69Q1NMX2qwAkzQSurSekiIjVc/plf6i8jg/tv3PlddStU4vg6b6JdAlFRIxcnVoEr5H0WDktijuLHyunbXuzyqOLiIjKdXpU5ag6A4mIiKHRzRATERExgiURREQ0XBJBRETDJRFERDRcEkFERMMlEURENFwSQUREwyURREQ0XBJBRETDJRFERDRcEkFERMMlEURENFwSQUREwyURREQ0XBJBRETDJRFERDRcEkFERMNVmggkTZN0q6Qlkk5qs34vSddLWinpkCpjiYiI9ipLBJJGAWcABwBTgMMlTelX7C5gOnBBVXFERERnnR5ev6amAktsLwWQdCFwELCor4DtZeW6ZyuMIyIiOqiya2gCsLxlfkW5bJVJOlrSAkkLent710pwERFRGBYni22fabvHds/48eOHOpyIiBGlykRwN7BNy/zEcllERKxDqkwE84GdJE2WNBo4DJhTYX0REbEaKksEtlcCM4B5wGJglu2Fkk6TdCCApNdLWgG8C/iOpIVVxRMREe1VedUQtucCc/stO7Vlej5Fl1FERAyRYXGyOCIiqpNEEBHRcEkEERENl0QQEdFwSQQREQ2XRBAR0XBJBBERDZdEEBHRcJXeULauOf2yP1Rex4f237nyOiIi1qZGJYKIpqj6oCcHPCNLuoYiIhouiSAiouGSCCIiGi6JICKi4XKyOCqVK7Ui1n1pEURENFwSQUREwyURREQ0XBJBRETDJRFERDRcrhqKES1DLUQMLi2CiIiGSyKIiGi4dA1FVCTdUjFcpEUQEdFwaRHUJEMtRIxsw/n/eFoEERENV2mLQNI04GvAKOBs25/vt/5FwPeA1wEPAe+2vazKmCKiWsP5yLipKmsRSBoFnAEcAEwBDpc0pV+xDwCP2N4ROB34t6riiYiI9qrsGpoKLLG91PZTwIXAQf3KHAScV07PBvaVpApjioiIfmS7mh1LhwDTbB9Zzr8X2N32jJYyN5dlVpTzt5dlHuy3r6OBo8vZXYBbKwm6vXHAg4OWSt2pO3Wn7nW77u1sj2+3YlhcNWT7TODMoahb0gLbPak7dafu1D1S6u6vyq6hu4FtWuYnlsvalpG0PrA5xUnjiIioSZWJYD6wk6TJkkYDhwFz+pWZA/xjOX0IcIWr6quKiIi2Kusasr1S0gxgHsXlo+fYXijpNGCB7TnATOD7kpYAD1Mki3XNkHRJpe7UnbpTd10qO1kcERHDQ+4sjohouCSCiIiGSyIYgKRpkm6VtETSSTXXfY6kB8r7LOqsdxtJv5C0SNJCSSfUWPeGkq6V9Puy7k/VVXdLDKMk/U7ST4eg7mWSbpJ0g6QFNdc9VtJsSbdIWizpDTXVu0v5fvtej0n6YB11l/V/qPy3drOkH0jasMa6TyjrXVjnex6Q7bz6vShObt8ObA+MBn4PTKmx/r2A3YCba37fWwG7ldNjgD/U9b4BAZuW0xsA1wB71Pz+TwQuAH5aZ71l3cuAcXXXW9Z9HnBkOT0aGDsEMYwC7qO46amO+iYAdwAblfOzgOk11f1K4GZgY4oLdi4HdhyK777vlRZBe90Mj1EZ2/9NcRVVrWzfa/v6cvpxYDHFf5g66rbtP5WzG5Sv2q5kkDQR+Afg7LrqXBdI2pziwGMmgO2nbD86BKHsC9xu+84a61wf2Ki8h2lj4J6a6v0b4Brbf7a9ErgSOLimuttKImhvArC8ZX4FNf0griskTQJeS3FkXledoyTdADwAXGa7trqBrwL/DDxbY52tDFwq6bpySJW6TAZ6ge+W3WJnS9qkxvr7HAb8oK7KbN8NfAm4C7gX+KPtS2uq/mbgzZK2lLQx8Pc8/+bb2iURxAtI2hT4IfBB24/VVa/tZ2zvSnEX+lRJr6yjXklvAx6wfV0d9Q3gTbZ3oxit9zhJe9VU7/oU3ZD/bvu1wBNA3efERgMHAhfVWOcWFK38ycDWwCaSjqijbtuLKUZavhS4BLgBeKaOugeSRNBeN8NjjEiSNqBIAv9h++KhiKHsmvgFMK2mKt8IHChpGUU34D6Szq+pbuC5I1RsPwD8iKJ7sg4rgBUtra/ZFImhTgcA19u+v8Y69wPusN1r+2ngYmDPuiq3PdP262zvBTxCcT5uyCQRtNfN8BgjTjkE+Exgse2v1Fz3eEljy+mNgP2BW+qo2/bJtifankTxXV9hu5ajQwBJm0ga0zcN/C1F90HlbN8HLJe0S7loX2BRHXW3OJwau4VKdwF7SNq4/He/L8U5sVpIekn5d1uK8wMX1FV3O8Ni9NG6eYDhMeqqX9IPgL2BcZJWAJ+0PbOGqt8IvBe4qeyrB/iY7bk11L0VcF75QKP1gFm2a7+Mc4i8FPhR+SiO9YELbF9SY/3HA/9RHvQsBd5fV8Vl4tsfOKauOgFsXyNpNnA9sBL4HfUO+fBDSVsCTwPHDdEJ+udkiImIiIZL11BERMMlEURENFwSQUREwyURREQ0XBJBRETDJRFERDRcEkFERMP9f9TbtPnh8O6mAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "objects = [str(x) for x in range(10)]\n",
    "y_pos = np.arange(len(objects))\n",
    "plt.bar(y_pos, probabilities, align='center', alpha=0.5)\n",
    "plt.xticks(y_pos, objects)\n",
    "plt.ylabel('Predicted Probability')\n",
    "plt.title('LFC-w1a1 Predictions for Image')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Call Brevitas FINN-ONNX export and visualize with Netron\n",
    "\n",
    "Brevitas comes with built-in FINN-ONNX export functionality. This is similar to the regular ONNX export capabilities of PyTorch, with a few differences:\n",
    "\n",
    "1. The weight quantization logic is not exported as part of the graph; rather, the quantized weights themselves are exported.\n",
    "2. Special quantization annotations are used to preserve the low-bit quantization information. ONNX (at the time of writing) supports 8-bit quantization as the minimum bitwidth, whereas FINN-ONNX quantization annotations can go down to binary/bipolar quantization.\n",
    "3. Low-bit quantized activation functions are exported as MultiThreshold operators.\n",
    "\n",
    "It's actually quite straightforward to export ONNX from our Brevitas model as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/brevitas_cnv_lfc/training_scripts/models/LFC.py:85: TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.\n",
      "  x = 2.0 * x - torch.tensor([1.0]).to(self.device)\n"
     ]
    }
   ],
   "source": [
    "import brevitas.onnx as bo\n",
    "export_onnx_path = \"/tmp/LFCW1A1.onnx\"\n",
    "input_shape = (1, 1, 28, 28)\n",
    "bo.export_finn_onnx(lfc, input_shape, export_onnx_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's examine what the exported ONNX model looks like. For this, we will use the Netron visualizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving '/tmp/LFCW1A1.onnx' at http://0.0.0.0:8081\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://0.0.0.0:8081/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7f3d330b6ac8>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "showInNetron('/tmp/LFCW1A1.onnx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When running this notebook in the FINN Docker container, you should be able to see an interactive visualization of the imported network above, and click on individual nodes to inspect their parameters. If you look at any of the MatMul nodes, you should be able to see that the weights are all {-1, +1} values, and the activations are Sign functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Import into FINN and call cleanup transformations\n",
    "\n",
    "We will now import this ONNX model into FINN using the ModelWrapper, and examine some of the graph attributes from Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "input: \"37\"\n",
       "input: \"38\"\n",
       "output: \"40\"\n",
       "op_type: \"MatMul\""
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from finn.core.modelwrapper import ModelWrapper\n",
    "model = ModelWrapper(export_onnx_path)\n",
    "model.graph.node[9]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ModelWrapper exposes a range of other useful functions as well. For instance, by convention the second input of the MatMul node will be a pre-initialized weight tensor, which we can view using the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-1., -1., -1., ..., -1., -1.,  1.],\n",
       "       [-1.,  1., -1., ...,  1., -1., -1.],\n",
       "       [ 1., -1.,  1., ..., -1., -1., -1.],\n",
       "       ...,\n",
       "       [ 1.,  1., -1., ...,  1.,  1.,  1.],\n",
       "       [-1., -1.,  1., ...,  1.,  1., -1.],\n",
       "       [ 1.,  1., -1., ...,  1., -1., -1.]], dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.get_initializer(model.graph.node[9].input[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also examine the quantization annotations and shapes of various tensors using the convenience functions provided by ModelWrapper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<DataType.BIPOLAR: 8>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.get_tensor_datatype(model.graph.node[9].input[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[784, 1024]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.get_tensor_shape(model.graph.node[9].input[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to operate further on this model in FINN, it is a good idea to execute certain \"cleanup\" transformations on this graph. Here, we will run shape inference and constant folding on this graph, and visualize the resulting graph in Netron again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from finn.transformation.fold_constants import FoldConstants\n",
    "from finn.transformation.infer_shapes import InferShapes\n",
    "model = model.transform(InferShapes())\n",
    "model = model.transform(FoldConstants())\n",
    "export_onnx_path_transformed = \"/tmp/LFCW1A1-clean.onnx\"\n",
    "model.save(export_onnx_path_transformed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Stopping http://0.0.0.0:8081\n",
      "Serving '/tmp/LFCW1A1-clean.onnx' at http://0.0.0.0:8081\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://0.0.0.0:8081/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7f3d3380aef0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "showInNetron('/tmp/LFCW1A1-clean.onnx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the resulting graph has become smaller and simpler. Specifically, the input reshaping is now a single Reshape node instead of the Shape -> Gather -> Unsqueeze -> Concat -> Reshape sequence. We can now use the internal ONNX execution capabilities of FINN to ensure that we still get the same output from this model as we did with PyTorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-1.5095654 , -2.915617  ,  0.764004  , -1.8118242 , -2.308991  ,\n",
       "        -2.6900144 , -1.520713  , -3.4965858 , -0.47711682, -2.9628415 ]],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import finn.core.onnx_exec as oxe\n",
    "input_dict = {\"0\": nph.to_array(input_tensor)}\n",
    "output_dict = oxe.execute_onnx(model, input_dict)\n",
    "produced_finn = output_dict[list(output_dict.keys())[0]]\n",
    "\n",
    "produced_finn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.isclose(produced, produced_finn).all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have succesfully verified that the transformed and cleaned-up FINN graph still produces the same output, and can now use this model for further processing in FINN."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
